import numpy as np
from sklearn.cluster import KMeans
import gradio as gr

def get_centers(img, k):
    # 将图像转换为二维数组
    img = img.reshape(-1, 3)

    # 应用 KMeans 聚类
    kmeans = KMeans(n_clusters=k)
    kmeans.fit(img)  # 训练模型
    
    # 获取聚类中心
    centers = kmeans.cluster_centers_
    return centers

def create_image(image):
    k = 5
    edge_length = 100
    centers = get_centers(image, k)
    
    # 创建一个掩模
    output_image = np.ones((100, edge_length*k, 3)) * 255  # 创建全白图像
    for i in range(k):
        output_image[:, i*edge_length:(i+1)*edge_length, :] = centers[i]  # 每个区域宽度为100
    
    # 确保数据类型为 uint8
    output_image = np.clip(output_image, 0, 255).astype(np.uint8)  # 限制范围并转换类型

    # 将 centers 转换为两位十六进制字符串
    vectorized_to_hex = np.vectorize(lambda x: f"{int(x):02x}")  # 转换为整数后再格式化
    centers_hex = vectorized_to_hex(centers)

    # 生成 colornums_list
    colornums_list = ['#'+''.join(centers_hex[i]) for i in range(len(centers_hex))]
    colornums_string = "\n".join(colornums_list)  # 每个元素换行

    return output_image, colornums_string  # 返回图像和 colornums_string

# 创建 Gradio 接口
iface = gr.Interface(
    fn=create_image,
    inputs=gr.Image(type="numpy"),
    outputs=[gr.Image(type='numpy'), gr.Textbox(label="Color Numbers")]  # 增加输出框
)

# 启动 Gradio 应用
iface.launch()